# set tensorflow backend for allowing GPU memory growth
# in order to avoid CUDNN_STATUS_INTERNAL_ERROR
import tensorflow as tf
from keras import backend as K
config = tf.ConfigProto(allow_soft_placement = True)
config.gpu_options.allow_growth = True
sess = tf.Session(config = config)
K.set_session(sess)

# we can also run the detector on CPU by disabling the GPU
# uncomment the follow two lines to enable CPU
# import os
# os.environ["CUDA_VISIBLE_DEVICES"]="-1"

# import packages
from keras.preprocessing.image import img_to_array
from keras.models import load_model
import numpy as np
import argparse
import imutils
import cv2

# construct the argument parser
ap = argparse.ArgumentParser()
ap.add_argument("-c", "--cascade", required = True,
    help = "path to where the face cascade resides")
ap.add_argument("-m", "--model", required = True,
    help = "path to pre-trained emotion detector CNN")
ap.add_argument("-v", "--video",
    help = "path to the (optional) video file")
args = vars(ap.parse_args())

# load the face detector cascade, emotion detection CNN, then define the
# list of emotion labels
detector = cv2.CascadeClassifier(args["cascade"])
model = load_model(args["model"])
EMOTIONS = ["angry", "scared", "happy", "sad", "surprised", "neutral"]

# if a video path is not provided, grab the reference to the webcam
if not args.get("video", False):
    camera = cv2.VideoCapture(0)

# otherwise, load the video
else:
    camera = cv2.VideoCapture(args["video"])

# keep looping
while True:
    # grab the current frame
    (grabbed, frame) = camera.read()

    # if we are viewing a video and we did not grab a frame
    # then we have reached the end of the video
    if args.get("video") and not grabbed:
        break

    # resize the frame and convert it to grayscale
    frame = imutils.resize(frame, width = 300)
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # initialize the canvas for the visualization
    # then clone the frame to draw on it
    canvas = np.zeros((220, 300, 3), dtype = "uint8")
    frameClone = frame.copy()

    # detect faces in the input frame
    rects = detector.detectMultiScale(gray, scaleFactor = 1.1,
        minNeighbors = 5, minSize = (30, 30), flags = cv2.CASCADE_SCALE_IMAGE)

    # ensure at least one face is found before continuing
    if len(rects) > 0:
        # determine the largest face area
        rect = sorted(rects, reverse = True,
            key = lambda x: (x[2] - x[1]) * (x[3] - x[1]))[0]
        (fX, fY, fW, fH) = rect

        # extract the face ROI from the image
        # then preprocess it for the network
        roi = gray[fY : fY + fH, fX : fX + fW]
        roi = cv2.resize(roi, (48, 48))
        roi = roi.astype("float") / 255.0
        roi = img_to_array(roi)
        roi = np.expand_dims(roi, axis = 0)

        # make a prediction on the ROI, then lookup the class label
        preds = model.predict(roi)[0]
        label = EMOTIONS[preds.argmax()]

        # loop over the labels + probabilities and draw them
        for (i, (emotion, prob)) in enumerate(zip(EMOTIONS, preds)):
            # construct the label text
            text = "{}: {:.2f}%".format(emotion, prob * 100)

            # draw the label + probability bar on canvas
            w = int(prob * 300)
            cv2.rectangle(canvas, (5, (i * 35) + 5), (w, (i * 35) + 35),
                (0, 0, 255), -1)
            cv2.putText(canvas, text, (10, (i * 35) + 23), cv2.FONT_HERSHEY_SIMPLEX,
                0.45, (255, 255, 255), 2)

        # draw the label on the frame
        cv2.putText(frameClone, label, (fX, fY - 10), cv2.FONT_HERSHEY_SIMPLEX,
            0.45, (0, 0, 255), 2)
        cv2.rectangle(frameClone, (fX, fY), (fX + fW, fY + fH),
            (0, 0, 255), 2)

    # show classifications + probabilities
    cv2.imshow("Face", frameClone)
    cv2.imshow("probabilities", canvas)

    # if the 'q' key is pressed, stop the loop
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break

# cleanup the camera and close any open windows
camera.release()
cv2.destroyAllWindows()
